# Copyright (c) 2006 Charl P. Botha <http://cpbotha.net/>
# All rights reserved.

# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions
# are met:
# 1. Redistributions of source code must retain the above copyright
#    notice, this list of conditions and the following disclaimer.
# 2. Redistributions in binary form must reproduce the above copyright
#    notice, this list of conditions and the following disclaimer in the
#    documentation and/or other materials provided with the distribution.
# 3. The name of the author may not be used to endorse or promote products
#    derived from this software without specific prior written permission.

# THIS SOFTWARE IS PROVIDED BY THE AUTHOR ``AS IS'' AND ANY EXPRESS OR
# IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES
# OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED.
# IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR ANY DIRECT, INDIRECT,
# INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT
# NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF
# THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


import pywintypes
import re
from utils import ci_str_cmp
import win32api
import win32con
import win32gui

CU_ENV_KEYSPEC = (win32con.HKEY_CURRENT_USER, "Environment")
SYS_ENV_KEYSPEC = \
                (win32con.HKEY_LOCAL_MACHINE,
                 "SYSTEM\\CurrentControlSet\\Control\\Session "
                 "Manager\\Environment")

class EnvKey:
    """Class encapsulating complete registry key, especially for handling
    the environment keys in CU_ENV_KEYSPEC and SYS_ENV_KEYSPEC.
    """

    def __init__(self, keyspec):
        self.keyspec = keyspec
        # dictionary mapping from name to (value, type) tuples
        self.items_dict = {}
        # decorator list containing names in original order
        self.names_list = []
        
        # dictionary similar to items_dict used for previewing
        # before applying to underlying key
        self.preview_items_dict = {}
        self.preview_names_list = []

    def apply_changes(self):
        """Apply items_dict to the registry that we encapsulate.

        Names that are in the registry, but not in our dict, will be deleted!
        We apply our changes in the least destructive way: first we call
        SetValue on all the known values we have in our items_dict, and only
        then do we delete the values that aren't amongst our known values.
        The only drawback of this is that the registry isn't explicitly
        sorted; that option would have been too high risk (delete all keys
        first, then write in sorted order).  However, by sorting after every
        read and after every parse of passed text, the names are always shown
        in sorted order.

        @raises pywintypes.error: if an error occurs writing to registry.
        """

        # open key
        hkey = self._open_envkey(self.keyspec)

        # remove all values from registry which are not in our new items_dict
        names_list, items_dict = self._extract_all_items(hkey)
        
        for name in names_list:
            if name not in self.items_dict:
                win32api.RegDeleteValue(hkey, name)
        
        # set all new values ################################################
        for name in self.names_list:
            # rt is (value, type)
            val, typ = self.items_dict[name]
            win32api.RegSetValueEx(hkey, name, 0, typ, val)

        
        # finally close the key
        self._close_key(hkey)

        # we have to send this message so that all programs which are
        # started up after our changes actually see them.
        # I found this in the cookbook at:
        # http://aspn.activestate.com/ASPN/Cookbook/Python/Recipe/55993
        # also see:
        # http://support.microsoft.com/?kbid=104011
        win32gui.SendMessageTimeout(
            win32con.HWND_BROADCAST, win32con.WM_SETTINGCHANGE, 0,
            'Environment', win32con.SMTO_ABORTIFHUNG, 100)        

    def _open_envkey(self, keyspec):
        """Returns opened key.

        @param keyspec: either CU_ENV_KEYSPEC or SYS_ENV_KEYSPEC, depending on
        what you require.
        """

        privileges = win32con.KEY_SET_VALUE | win32con.KEY_READ
        hkey = win32api.RegOpenKey(keyspec[0], keyspec[1], 0, privileges)

        return hkey

    def _close_key(self, hkey):
        win32api.RegCloseKey(hkey)

    def export_to_reg_file(self, filename):
        """Saves registry to REG file.

        For this to work, the process has to run with SE_BACKUP_NAME
        privileges.

        @raises pywintypes.error: if export to reg fails.
        """
        
        hkey = self._open_envkey(self.keyspec)
        win32api.RegSaveKey(hkey, filename)
        self._close_key(hkey)

    def _extract_all_items(self, hkey, uppercase_names=False):
        """Extract all items from hkey.

        @param uppercase_names: If True, all names will be uppercased.  This
        functionality is not currently being used.
        @returns: tuple with (list of names, dictionary).  dictionary keyed on
        item name, value is a tuple with variable value and type
        (win32con.REG_EXPAND_SZ or win32con.REG_SZ).
        """

        items_dict = {}
        names_list = []
        i = 0
        try:
            while True:
                # rt is a three element tuple: name, value, type (REG_SZ or
                # REG_EXPAND_SZ)
                rt = win32api.RegEnumValue(hkey, i)
                if rt[2] == win32con.REG_SZ or \
                   rt[2] == win32con.REG_EXPAND_SZ:
                    name = rt[0]

                    if uppercase_names:
                        name = name.upper()
                        
                    names_list.append(name)
                    items_dict[name] = (rt[1], rt[2])
                
                i += 1

        except pywintypes.error:
            # we get this error when we've reached the last element
            pass

        # the windows environment editor always explicitly sorts env
        # variables... we don't, but we do display them in sorted order.
        # we sort here, and also right after the parsing step in set_text()
        names_list.sort(ci_str_cmp)
        
        return names_list, items_dict

    def sync(self):
        try:
            hkey = self._open_envkey(self.keyspec)
            self.names_list, self.items_dict = self._extract_all_items(hkey)
            self._close_key(hkey)
        except pywintypes.error:
            raise

    def get_text(self, preview=False):
        """Return all names in current key as text.

        @param preview: return text for preview dict and not the current items
        dict.
        """

        text_list = []

        if preview:
            nlist = self.preview_names_list
            idict = self.preview_items_dict
        else:
            nlist = self.names_list
            idict = self.items_dict
            
        for name in nlist:
            text_list.append('%s = %s' % (name, idict[name][0]))

        new_text = '\n\n'.join(text_list)

        return new_text

    def set_text(self, text, preview=False):
        """Parse and apply text to items dict.

        This will not make any changes to the underlying registry.

        @param preview: apply to preview_dict and not to real
        items_dict.
        """

        # parse text into VARIABLE = value pairs, determine type as well.
        # \w - any alphanumeric character
        # \s - any whitespace
        # usually, this expression won't match past end of line
        po = re.compile('^\s*(\w+)\s*=\s*(.*)$')
        po_expand = re.compile('.*%.*%.*')

        new_dict = {}
        new_names_list = []
        # split the text on newlines
        txt_lines = text.split('\n')
        # now go through the lines getting out the name, value pairs
        for txt_line in txt_lines:
            mo = po.search(txt_line)
            if mo:
                if mo.group(1) not in new_dict:
                    # determine whether we should use REG_SZ or REG_EXPAND_SZ
                    # if the value contains %blabla% anywhere, it's EXPAND
                    if po_expand.match(mo.group(2)):
                        regtype = win32con.REG_EXPAND_SZ
                    else:
                        regtype = win32con.REG_SZ
                    
                    new_names_list.append(mo.group(1))
                    new_dict[mo.group(1)] = (mo.group(2), regtype)

                    # if the user has screwed up and made two equal names
                    # we only store the first

        # sort here as well as when we extract all items from the registry
        new_names_list.sort(ci_str_cmp)

        if preview:
            self.preview_items_dict = new_dict
            self.preview_names_list = new_names_list

        else:
            self.items_dict = new_dict
            self.names_list = new_names_list
